import pymysql
import json
import os

def cleanup_test_data():
    # 数据库连接配置
    config = {
        'host': '127.0.0.1',
        'port': 3306,
        'user': 'root',
        'password': '123456',
        'database': 'drcurd',
        'charset': 'utf8'
    }
    
    # 读取测试数据中的用户名
    data_file = "../data/register_data.json"
    with open(data_file, 'r', encoding='utf-8') as f:
        test_data = json.load(f)
    
    # 提取所有测试用户名
    test_usernames = [item['req_data']['username'] for item in test_data]
    
    try:
        # 连接数据库
        conn = pymysql.connect(**config)
        cursor = conn.cursor()
        
        # 删除测试数据
        for username in test_usernames:
            sql = "DELETE FROM tshop.user WHERE username = %s"
            cursor.execute(sql, (username,))
            if cursor.rowcount > 0:
                print(f"已删除用户: {username}")
            else:
                print(f"用户不存在: {username}")
        
        # 提交事务
        conn.commit()
        print(f"清理完成，共处理 {len(test_usernames)} 个测试用户名")
        
    except Exception as e:
        print(f"清理失败: {e}")
        if 'conn' in locals():
            conn.rollback()
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

if __name__ == "__main__":
    cleanup_test_data() 